﻿/**
 * TCP服务器，使用epoll实现大并发监听和接收数据
 * @author: wangjian@ruiboyun.net
 */
#include <iostream>
#include <string.h> // memset
#include <unistd.h> // close
#include <sys/socket.h> // socket, bind, listen, accept
#include <netinet/in.h> // sockaddr_in
#include <arpa/inet.h> // inet_addr
#include <fcntl.h> // fcntl
#include <errno.h> // errno
#include <sys/epoll.h> // epoll
#include <unordered_map>
#include <string>
#include <vector>

#include "connections.hpp"

#define MAX_EVENTS 3000

class TCPServer {
public:
    TCPServer(int port,Connections& conns) : listenSocket(-1), port(port),connections(conns),msgCount(0) {}

    ~TCPServer() {
        if (listenSocket != -1)
            close(listenSocket);
    }

    void start() {
        initSocket();
        setupEpoll();
        listen(listenSocket, SOMAXCONN);
        std::cout << "Server started on port " << port << std::endl;
        run();
    }

    size_t getConnectionCount() {
        return connections.size();
    }

private:
    int listenSocket;
    int epollFd;
    int port;
    int msgCount;
    Connections connections;

    void initSocket() {
        listenSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (listenSocket == -1) {
            perror("Error creating socket");
            exit(EXIT_FAILURE);
        }

        int optval = 1;
        setsockopt(listenSocket, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));

        struct sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_port = htons(port);
        addr.sin_addr.s_addr = htonl(INADDR_ANY);

        if (bind(listenSocket, (struct sockaddr*)&addr, sizeof(addr)) == -1) {
            perror("Error binding socket");
            close(listenSocket);
            exit(EXIT_FAILURE);
        }
    }

    void setupEpoll() {
        epollFd = epoll_create1(0);
        if (epollFd == -1) {
            perror("Error creating epoll instance");
            close(listenSocket);
            exit(EXIT_FAILURE);
        }

        struct epoll_event ev;
        ev.events = EPOLLIN | EPOLLET;
        ev.data.fd = listenSocket;
        if (epoll_ctl(epollFd, EPOLL_CTL_ADD, listenSocket, &ev) == -1) {
            perror("Error adding listen socket to epoll");
            close(listenSocket);
            close(epollFd);
            exit(EXIT_FAILURE);
        }

        fcntl(listenSocket, F_SETFL, O_NONBLOCK);
    }

    void run() {
        while (true) {
            std::vector<struct epoll_event> events(MAX_EVENTS);
            int numEvents = epoll_wait(epollFd, events.data(), MAX_EVENTS, -1);
            if (numEvents == -1) {
                perror("Error in epoll_wait");
                break;
            }

            for (int i = 0; i < numEvents; ++i) {
                if (events[i].data.fd == listenSocket) {
                    handleNewConnection();
                }
                else {
                    handleClient(events[i].data.fd);
                }
            }
        }
    }

    /**
     * 接收新连接
     */
    void handleNewConnection() {
        int clientSocket = accept(listenSocket, nullptr, nullptr);
        if (clientSocket == -1) {
            if (errno != EAGAIN && errno != EWOULDBLOCK) {
                perror("Error accepting new connection");
            }
            return;
        }

        //添加到epoll中
        struct epoll_event ev;
        ev.events = EPOLLIN | EPOLLET;
        ev.data.fd = clientSocket;
        if (epoll_ctl(epollFd, EPOLL_CTL_ADD, clientSocket, &ev) == -1) {
            perror("Error adding client socket to epoll");
            close(clientSocket);
            return;
        }

        //设置非阻塞模式
        fcntl(clientSocket, F_SETFL, O_NONBLOCK);

        Client client(clientSocket);
        std::string key = std::to_string(clientSocket);  // Use socket number as key for simplicity
        connections.insert(key,client);
        std::cout << "Client connected " << client.address <<  "  Total " << std::to_string(connections.size()) << std::endl;
    }

    /**
     * 处理客户端消息,交给客户端对象去处理
     * 如果客户端receiveData返回结果为0，则关闭连接
     */
    void handleClient(int clientSocket) {
        std::string key = std::to_string(clientSocket);
        Client client = connections.get(key);
        if(!client.clientSocket){
            std::cout << "Unhandled message !" << std::endl;
            return;
        }

        //debug
        // ++msgCount;
        // if(msgCount % 1000 <= 1){
        //     std::cout << "Connections:" << std::to_string(connections.size()) << "  messages:" << std::to_string(++msgCount) << std::endl;
        // }
        
        //handle error
        if (client.receiveData(clientSocket) == 0) {
            // Client disconnected
            close(clientSocket);
            epoll_ctl(epollFd, EPOLL_CTL_DEL, clientSocket, nullptr);
            
            int r = connections.erase(key);
            std::cout << "Client disconnected " << " Total " << std::to_string(connections.size()) << std::endl;
        }

        
    }
};
